# SAC强化学习算法
# src/algorithms/sac.py
# src/algorithms/sac.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from typing import Dict, List, Tuple, Any
from collections import deque
import random

class Actor(nn.Module):
    """演员网络"""
    
    def __init__(self, state_dim: int, action_dim: int, max_action: float = 1.0):
        super(Actor, self).__init__()
        self.max_action = max_action
        
        self.fc1 = nn.Linear(state_dim, 256)
        self.fc2 = nn.Linear(256, 256)
        self.mean = nn.Linear(256, action_dim)
        self.log_std = nn.Linear(256, action_dim)
        
        # 限制log_std的范围
        self.log_std_min = -20
        self.log_std_max = 2
    
    def forward(self, state):
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        
        mean = self.mean(x)
        log_std = self.log_std(x)
        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)
        
        return mean, log_std
    
    def sample(self, state):
        mean, log_std = self.forward(state)
        std = torch.exp(log_std)
        
        # 重参数化采样
        normal = torch.randn_like(mean)
        x_t = mean + std * normal
        
        # 使用tanh压缩到[-1,1]
        action = torch.tanh(x_t)
        
        # 计算对数概率
        log_prob = torch.distributions.Normal(mean, std).log_prob(x_t)
        log_prob -= torch.log(1 - action.pow(2) + 1e-6)
        log_prob = log_prob.sum(1, keepdim=True)
        
        return action * self.max_action, log_prob
    
    def get_action(self, state):
        """获取确定性动作（用于测试）"""
        mean, _ = self.forward(state)
        return torch.tanh(mean) * self.max_action

class Critic(nn.Module):
    """评论家网络"""
    
    def __init__(self, state_dim: int, action_dim: int):
        super(Critic, self).__init__()
        
        # Q1网络
        self.q1_fc1 = nn.Linear(state_dim + action_dim, 256)
        self.q1_fc2 = nn.Linear(256, 256)
        self.q1_out = nn.Linear(256, 1)
        
        # Q2网络
        self.q2_fc1 = nn.Linear(state_dim + action_dim, 256)
        self.q2_fc2 = nn.Linear(256, 256)
        self.q2_out = nn.Linear(256, 1)
    
    def forward(self, state, action):
        sa = torch.cat([state, action], dim=1)
        
        # Q1
        q1 = F.relu(self.q1_fc1(sa))
        q1 = F.relu(self.q1_fc2(q1))
        q1 = self.q1_out(q1)
        
        # Q2
        q2 = F.relu(self.q2_fc1(sa))
        q2 = F.relu(self.q2_fc2(q2))
        q2 = self.q2_out(q2)
        
        return q1, q2

class SACPortfolioOptimizer:
    """SAC投资组合优化器"""
    
    def __init__(self, state_dim: int, action_dim: int, 
                 lr: float = 3e-4, alpha: float = 0.2, 
                 gamma: float = 0.99, tau: float = 0.005,
                 buffer_size: int = 1000000, batch_size: int = 256):
        
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.gamma = gamma
        self.tau = tau
        self.alpha = alpha
        self.batch_size = batch_size
        
        # 网络初始化
        self.actor = Actor(state_dim, action_dim).to(self.device)
        self.critic = Critic(state_dim, action_dim).to(self.device)
        self.critic_target = Critic(state_dim, action_dim).to(self.device)
        
        # 复制参数到目标网络
        self.critic_target.load_state_dict(self.critic.state_dict())
        
        # 优化器
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr)
        
        # 经验回放缓冲区
        self.replay_buffer = deque(maxlen=buffer_size)
        
        # 自动熵调整
        self.target_entropy = -action_dim
        self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
        self.alpha_optimizer = optim.Adam([self.log_alpha], lr=lr)
    
    def get_state(self, portfolio_weights: np.ndarray, 
                  market_data: Dict[str, Any]) -> np.ndarray:
        """构建状态向量"""
        # 当前组合权重
        state = list(portfolio_weights)
        
        # 市场数据
        state.extend([
            market_data.get('vix', 0),  # 恐慌指数
            market_data.get('market_return', 0),  # 市场收益
            market_data.get('market_volatility', 0),  # 市场波动率
        ])
        
        # 技术指标
        state.extend([
            market_data.get('rsi', 50),  # RSI
            market_data.get('macd', 0),  # MACD
            market_data.get('volume_ratio', 1),  # 成交量比率
        ])
        
        return np.array(state, dtype=np.float32)
    
    def calculate_reward(self, prev_weights: np.ndarray, 
                        new_weights: np.ndarray,
                        returns: np.ndarray, 
                        transaction_cost: float = 0.005) -> float:
        """计算奖励函数"""
        # 组合收益
        portfolio_return = np.dot(new_weights, returns)
        
        # 交易成本
        turnover = np.sum(np.abs(new_weights - prev_weights))
        cost = transaction_cost * turnover
        
        # 风险调整（使用预期收益/波动率作为Sharpe比率的近似）
        portfolio_vol = np.sqrt(np.dot(new_weights**2, returns**2))  # 简化版本
        risk_adjusted_return = portfolio_return / (portfolio_vol + 1e-6)
        
        # 约束惩罚
        constraint_penalty = 0
        # 单一资产权重不超过20%
        if np.max(new_weights) > 0.2:
            constraint_penalty += 10 * (np.max(new_weights) - 0.2)
        
        # 组合总权重约束
        weight_sum_penalty = 10 * abs(np.sum(new_weights) - 1.0)
        
        reward = risk_adjusted_return - cost - constraint_penalty - weight_sum_penalty
        
        return reward
    
    def store_transition(self, state: np.ndarray, action: np.ndarray,
                        reward: float, next_state: np.ndarray, done: bool):
        """存储转移样本"""
        self.replay_buffer.append((state, action, reward, next_state, done))
    
    def sample_batch(self) -> Tuple[torch.Tensor, ...]:
        """采样批次数据"""
        batch = random.sample(self.replay_buffer, self.batch_size)
        
        state = torch.FloatTensor(np.array([e[0] for e in batch])).to(self.device)
        action = torch.FloatTensor(np.array([e[1] for e in batch])).to(self.device)
        reward = torch.FloatTensor(np.array([e[2] for e in batch])).unsqueeze(1).to(self.device)
        next_state = torch.FloatTensor(np.array([e[3] for e in batch])).to(self.device)
        done = torch.BoolTensor(np.array([e[4] for e in batch])).unsqueeze(1).to(self.device)
        
        return state, action, reward, next_state, done
    
    def update_networks(self):
        """更新网络"""
        if len(self.replay_buffer) < self.batch_size:
            return
        
        state, action, reward, next_state, done = self.sample_batch()
        
        # 更新Critic
        with torch.no_grad():
            next_action, next_log_prob = self.actor.sample(next_state)
            q1_next, q2_next = self.critic_target(next_state, next_action)
            q_next = torch.min(q1_next, q2_next) - self.alpha * next_log_prob
            q_target = reward + self.gamma * q_next * (~done)
        
        q1_current, q2_current = self.critic(state, action)
        critic_loss = F.mse_loss(q1_current, q_target) + F.mse_loss(q2_current, q_target)
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        
        # 更新Actor
        new_action, log_prob = self.actor.sample(state)
        q1_new, q2_new = self.critic(state, new_action)
        q_new = torch.min(q1_new, q2_new)
        
        actor_loss = (self.alpha * log_prob - q_new).mean()
        
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()
        
        # 更新alpha
        alpha_loss = -(self.log_alpha * (log_prob + self.target_entropy).detach()).mean()
        
        self.alpha_optimizer.zero_grad()
        alpha_loss.backward()
        self.alpha_optimizer.step()
        
        self.alpha = self.log_alpha.exp()
        
        # 软更新目标网络
        self.soft_update(self.critic_target, self.critic)
    
    def soft_update(self, target, source):
        """软更新目标网络"""
        for target_param, param in zip(target.parameters(), source.parameters()):
            target_param.data.copy_(target_param.data * (1.0 - self.tau) + param.data * self.tau)
    
    def get_action(self, state: np.ndarray, deterministic: bool = False) -> np.ndarray:
        """获取动作（权重调整）"""
        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        
        if deterministic:
            action = self.actor.get_action(state)
        else:
            action, _ = self.actor.sample(state)
        
        return action.cpu().data.numpy().flatten()
    
    def normalize_weights(self, weights: np.ndarray) -> np.ndarray:
        """归一化权重"""
        weights = np.maximum(weights, 0)  # 确保非负
        if np.sum(weights) > 0:
            weights = weights / np.sum(weights)
        return weights
    
    def save_model(self, filepath: str):
        """保存模型"""
        torch.save({
            'actor_state_dict': self.actor.state_dict(),
            'critic_state_dict': self.critic.state_dict(),
            'actor_optimizer': self.actor_optimizer.state_dict(),
            'critic_optimizer': self.critic_optimizer.state_dict(),
            'alpha': self.alpha,
            'log_alpha': self.log_alpha
        }, filepath)
    
    def load_model(self, filepath: str):
        """加载模型"""
        checkpoint = torch.load(filepath, map_location=self.device)
        
        self.actor.load_state_dict(checkpoint['actor_state_dict'])
        self.critic.load_state_dict(checkpoint['critic_state_dict'])
        self.actor_optimizer.load_state_dict(checkpoint['actor_optimizer'])
        self.critic_optimizer.load_state_dict(checkpoint['critic_optimizer'])
        self.alpha = checkpoint['alpha']
        self.log_alpha = checkpoint['log_alpha']
